#!/usr/bin/env python3
"""
Prometheus2 MCP Server

面向 Prometheus2 监控系统的 MCP Server，提供命令行工具和 API 调用的封装。

MCP tools 列表：
1. prometheus_version - 获取 Prometheus 版本信息
2. prometheus_help - 获取 Prometheus 帮助信息
3. prometheus_start_with_config - 使用指定配置文件启动 Prometheus
4. prometheus_start_with_listen_address - 使用指定监听地址启动 Prometheus
5. prometheus_management_api - 调用 Prometheus 管理 API
6. promtool_check_metrics - 使用 promtool 检查 metrics 格式
7. promtool_query_instant - 使用 promtool 执行即时查询
8. promtool_query_range - 使用 promtool 执行范围查询
9. prometheus_query_api - 调用 Prometheus 查询 API
10. tsdb_bench_write - 使用 tsdb 进行基准写入测试
11. tsdb_list_blocks - 使用 tsdb 列出数据块
12. tsdb_analyze - 使用 tsdb 分析数据块
"""

import subprocess
import json
from typing import List, Optional
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("Prometheus2 MCP Server")

@mcp.tool()
def prometheus_version() -> dict:
    """
    获取 Prometheus 版本信息
    
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["prometheus", "--version"],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": "prometheus --version",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "prometheus --version",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def prometheus_help() -> dict:
    """
    获取 Prometheus 帮助信息
    
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["prometheus", "--help"],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": "prometheus --help",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "prometheus --help",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def prometheus_start_with_config(config_file: str) -> dict:
    """
    使用指定配置文件启动 Prometheus
    
    参数:
        config_file: 配置文件路径
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["prometheus", f"--config.file={config_file}"],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": f"prometheus --config.file={config_file}",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"prometheus --config.file={config_file}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def prometheus_start_with_listen_address(address: str) -> dict:
    """
    使用指定监听地址启动 Prometheus
    
    参数:
        address: 监听地址，格式为 "IP:PORT"
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["prometheus", f"--web.listen-address={address}"],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": f"prometheus --web.listen-address={address}",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"prometheus --web.listen-address={address}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def prometheus_management_api(api_endpoint: str, method: str = "GET", data: Optional[str] = None) -> dict:
    """
    调用 Prometheus 管理 API
    
    参数:
        api_endpoint: API 端点，如 "/-/healthy", "/-/ready", "/-/reload", "/-/quit"
        method: HTTP 方法，GET 或 POST
        data: POST 请求的数据
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        curl_cmd = ["curl", "-s", f"http://localhost:9090{api_endpoint}"]
        if method.upper() == "POST":
            curl_cmd.extend(["-X", "POST"])
        if data:
            curl_cmd.extend(["--data", data])
            
        result = subprocess.run(
            curl_cmd,
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": " ".join(curl_cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"curl {method} http://localhost:9090{api_endpoint}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def promtool_check_metrics(metrics_data: str) -> dict:
    """
    使用 promtool 检查 metrics 格式
    
    参数:
        metrics_data: metrics 数据字符串
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["promtool", "check", "metrics"],
            input=metrics_data,
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": "promtool check metrics",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": "promtool check metrics",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def promtool_query_instant(server_url: str, query: str) -> dict:
    """
    使用 promtool 执行即时查询
    
    参数:
        server_url: Prometheus 服务器 URL
        query: 查询表达式
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["promtool", "query", "instant", server_url, query],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": f"promtool query instant {server_url} '{query}'",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"promtool query instant {server_url} '{query}'",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def promtool_query_range(server_url: str, query: str, start: str, end: str, step: str) -> dict:
    """
    使用 promtool 执行范围查询
    
    参数:
        server_url: Prometheus 服务器 URL
        query: 查询表达式
        start: 开始时间戳
        end: 结束时间戳
        step: 步长
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        result = subprocess.run(
            ["promtool", "query", "range", "--start", start, "--end", end, "--step", step, server_url, query],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": f"promtool query range --start {start} --end {end} --step {step} {server_url} '{query}'",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"promtool query range --start {start} --end {end} --step {step} {server_url} '{query}'",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def prometheus_query_api(api_endpoint: str, params: Optional[str] = None) -> dict:
    """
    调用 Prometheus 查询 API
    
    参数:
        api_endpoint: API 端点，如 "/api/v1/query", "/api/v1/series" 等
        params: 查询参数
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        url = f"http://localhost:9090{api_endpoint}"
        if params:
            url = f"{url}?{params}"
            
        result = subprocess.run(
            ["curl", "-s", url],
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": f"curl -s {url}",
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"curl -s http://localhost:9090{api_endpoint}",
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def tsdb_bench_write(input_file: str, output_dir: Optional[str] = None, metrics_count: Optional[int] = None) -> dict:
    """
    使用 tsdb 进行基准写入测试
    
    参数:
        input_file: 输入数据文件路径
        output_dir: 输出目录路径
        metrics_count: metrics 数量限制
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        cmd = ["tsdb", "bench", "write", input_file]
        if output_dir:
            cmd.extend(["--out", output_dir])
        if metrics_count:
            cmd.extend(["--metrics", str(metrics_count)])
            
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=60
        )
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def tsdb_list_blocks(db_path: Optional[str] = None, human_readable: bool = False) -> dict:
    """
    使用 tsdb 列出数据块
    
    参数:
        db_path: 数据库路径，如果为 None 则使用默认路径
        human_readable: 是否使用人类可读格式
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        cmd = ["tsdb", "ls"]
        if human_readable:
            cmd.append("-h")
        if db_path:
            cmd.append(db_path)
            
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def tsdb_analyze(db_path: Optional[str] = None, block_id: Optional[str] = None, limit: Optional[int] = None) -> dict:
    """
    使用 tsdb 分析数据块
    
    参数:
        db_path: 数据库路径，如果为 None 则使用默认路径
        block_id: 特定数据块 ID
        limit: 分析数量限制
        
    返回:
        dict: 包含执行结果的 JSON 结构
    """
    try:
        cmd = ["tsdb", "analyze"]
        if limit:
            cmd.extend(["--limit", str(limit)])
        if db_path:
            cmd.append(db_path)
        if block_id:
            cmd.append(block_id)
            
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30
        )
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()